package com.mycat.codegenerator.db;

import com.mycat.codegenerator.config.DBConfig;
import com.mycat.codegenerator.entity.Column;

import java.sql.*;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 解析数据库表信息处理程序
 * @author roger
 */
public class DBHandler {

    public static void main(String[] args) {

        List<String> tableNames = new ArrayList<String>() {{
            this.add("cg_code_generation");
            this.add("sys_ref_admin_role");
        }};

        DBHandler.getTables(tableNames);
        DBHandler.getTables();
    }


    /**
     * 获取数据库所有表信息
     *
     * @return   key->表名，value，表的所有字段信息
     */
    public static Map<String, List<Column>> getTables() {
        Map<String, List<Column>> map = new HashMap<>();
        Connection connection = null;
        ResultSet allTableResultSet = null;
        ResultSet tableResultSet = null;
        try {
            Class.forName(DBConfig.driver);
            // 获取链接
            connection = DriverManager.getConnection(DBConfig.url, DBConfig.user, DBConfig.password);
            // 获取此 Connection 对象的当前目录名称
            String currentCatalog = connection.getCatalog();
            //  获取一个 DatabaseMetaData 对象，该对象包含关于此 Connection 对象所连接的数据库的元数据。
            DatabaseMetaData databaseMetaData = connection.getMetaData();
            // 获取当前数据库所有表信息
            allTableResultSet = databaseMetaData.getTables(currentCatalog, null, null, new String[]{"TABLE"});

            while (allTableResultSet.next()) {
                // 获取表名
                String tableName = allTableResultSet.getString("TABLE_NAME");
                // 获取当前表信息
                tableResultSet = databaseMetaData.getColumns(currentCatalog, null, tableName, null);

                List<Column> columns = getColumns(tableResultSet);
                map.put(tableName, columns);

                System.out.println(tableName + " --> " + columns);
            }
            return map;
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (null != tableResultSet) {
                try {
                    tableResultSet.close();
                } catch (SQLException e) {
                    e.printStackTrace();
                }
            }
            if (null != allTableResultSet) {
                try {
                    allTableResultSet.close();
                } catch (SQLException e) {
                    e.printStackTrace();
                }
            }
            if (null != connection) {
                try {
                    connection.close();
                } catch (SQLException e) {
                    e.printStackTrace();
                }
            }
        }
        return map;
    }

    /**
     * 根据传入表名集合获取数据库所有表
     *
     * @return key->表名，value，表的所有字段信息
     */
    public static Map<String, List<Column>> getTables(List<String> tableNames) {
        Map<String, List<Column>> map = new HashMap<>();
        if (null == tableNames || tableNames.size() == 0) {
            return map;
        }
        Connection connection = null;
        ResultSet resultSet = null;
        try {
            Class.forName(DBConfig.driver);
            // 获取链接
            connection = DriverManager.getConnection(DBConfig.url, DBConfig.user, DBConfig.password);
            // 获取此 Connection 对象的当前目录名称
            String currentCatalog = connection.getCatalog();
            //  获取一个 DatabaseMetaData 对象，该对象包含关于此 Connection 对象所连接的数据库的元数据。
            DatabaseMetaData databaseMetaData = connection.getMetaData();

            for (String tableName : tableNames) {
                // 获取当前表信息
                resultSet = databaseMetaData.getColumns(currentCatalog, null, tableName, null);

                List<Column> columns = getColumns(resultSet);
                map.put(tableName, columns);

                System.out.println(tableName + " --> " + columns);
            }
            return map;
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (null != resultSet) {
                try {
                    resultSet.close();
                } catch (SQLException e) {
                    e.printStackTrace();
                }
            }
            if (null != connection) {
                try {
                    connection.close();
                } catch (SQLException e) {
                    e.printStackTrace();
                }
            }
        }
        return map;
    }

    /**
     * 解析结果集
     *
     * @param resultSet 当前表的结果集
     * @return 解析后的结果集
     * @throws SQLException
     */
    public static List<Column> getColumns(ResultSet resultSet) throws SQLException {
        List<Column> list = new ArrayList<>();
        if(null == resultSet){
            return list;
        }
        while (resultSet.next()) {
            Column column = new Column();
            column.setName(resultSet.getString("COLUMN_NAME"));
            column.setDataType(resultSet.getString("DATA_TYPE"));
            column.setIsAutoincrement(resultSet.getString("IS_AUTOINCREMENT"));
            column.setRemarks(resultSet.getString("REMARKS"));
            list.add(column);
        }
        return list;
    }
}
